#!/usr/bin/python

"""
This program draws gene clusters to make publication quality figures. It takes in
Genbank file, COG category file, and expects start and stop coordinates of region
to inspect.

Usage: dissertation_CompareGenes.py seq1.gbk seq2.gbk seq2.ptt cogs.t.list 
    <seq1start> <seq1stop> <seq2start> <seq2stop>
Examples:

dissertation_CompareGenes.py ../annotation/GKIL.v6.gbf NC_005125.1.gbk 
    NC_005125.ptt orthologs/orthomcl/cogs.t.list 10000 20000 10000 20000
dissertation_CompareGenes.py ../annotation/GKIL.v6.gbf NC_005125.1.gbk 
    NC_005125.ptt orthologs/orthomcl/cogs.t.list 773000 783000  2814800 2824800

dissertation_CompareGenes.py ../annotation/GKIL.v6.gbf NC_005125.1.gbk 
    NC_005125.ptt orthologs/orthomcl/cogs.t.list 2635000 2655000   
    178500 198500 (This one is comparing rhodopsin gene cluster between 
    GVIO and GKIL)
dissertation_CompareGenes.py ../annotation/GKIL.v6.gbf NC_005125.1.gbk 
    NC_005125.ptt orthologs/orthomcl/cogs.t.list 2625000 2665000 168500 208500

Note: Resolution is best if the segment in view is less than 10000bp.

Author: Jimmy Saw
Date of last update: 05-01-2012

"""

import sys
import re
import matplotlib.pyplot as plt
import pylab
import matplotlib
from matplotlib import mpl
from matplotlib.patches import Rectangle
from matplotlib.transforms import Bbox
from Bio import SeqIO
from Bio.SeqUtils import GC
import matplotlib.patches as mpatch

#Regex and other stuffs
cogcat = re.compile('\[(.*)\]\t(\w+)\t.*')
pttcog = re.compile('(COG\d{4})(\w).*')

cogdict = {
    'J' : '#2B60DE',
    'A' : '#F6358A',
    'K' : '#B048B5',
    'L' : '#8E35EF',
    'B' : '#D16587',
    'D' : '#C38EC7',
    'Y' : '#52F3FF',
    'V' : '#3EA99F',
    'T' : '#254117',
    'M' : '#41A317',
    'N' : '#00FF00',
    'Z' : '#FFFF00',
    'W' : '#FDD017',
    'U' : '#F88017',
    'O' : '#F660AB',
    'C' : '#FF0000',
    'G' : '#FAAFBA',
    'E' : '#7F5A58',
    'F' : '#C8B560',
    'H' : '#8B7500',
    'I' : '#C12869',
    'P' : '#57E964',
    'Q' : '#BCE954',
    'R' : '#F87431',
    'S' : '#ADA96E',
    '-' : '#D3D3D3'
    }

g1starts = []
g1stops = []
g2starts = []
g2stops = []
toplot = []

g1spanx1 = int(sys.argv[5])
g1spanx2 = int(sys.argv[6])
g2spanx1 = int(sys.argv[7])
g2spanx2 = int(sys.argv[8])

##Genome one Genbank file
g1seq = SeqIO.read(sys.argv[1], "gb")
g1length = len(g1seq.seq)
g1featdict = {}
for feat in g1seq.features:
    g1start = feat.location._start.position
    g1stop = feat.location._end.position
    g1starts.append(g1start)
    g1stops.append(g1stop)
    if g1spanx1 <= g1start <= g1spanx2 and g1spanx1 <= g1stop <= g1spanx2:
        if feat.type == 'CDS':
            g1featdict[feat.qualifiers['locus_tag'][0]] = feat
        if feat.type == 'tRNA':
            g1featdict[feat.qualifiers['locus_tag'][0]] = feat
        if feat.type == 'rRNA':
            g1featdict[feat.qualifiers['locus_tag'][0]] = feat

##Genome 2 Genbank file
g2seq = SeqIO.read(sys.argv[2], "gb")
g2length = len(g2seq.seq)
g2featdict = {}
for feat in g2seq.features:
    g2start = feat.location._start.position
    g2stop = feat.location._end.position
    g2starts.append(g2start)
    g2stops.append(g2stop)
    if g2spanx1 <= g2start <= g2spanx2 and g2spanx1 <= g2stop <= g2spanx2:
        if feat.type == 'CDS':
            g2featdict[feat.qualifiers['locus_tag'][0]] = feat
        if feat.type == 'tRNA':
            g2featdict[feat.qualifiers['locus_tag'][0]] = feat
        if feat.type == 'rRNA':
            g2featdict[feat.qualifiers['locus_tag'][0]] = feat

g1starts.sort()
g2stops.sort()
g2starts.sort()
g2stops.sort()

##Genome 2 ptt file
g2cogs = {}
g2pttfile = open(sys.argv[3], "rU")
g2ptt = g2pttfile.readlines()
for line in g2ptt[3:]:
    c = line.split('\t')
    g2ltag = c[5]
    g2gene = c[4]
    g2cog = '-'
    if pttcog.match(c[7]):
        p = pttcog.match(c[7])
        g2cog = p.group(1)
    g2cogs[g2ltag] = g2cog

cogcatfile = open(sys.argv[4], "rU")
cfl = cogcatfile.readlines()

cogcatdict = {}

for line in cfl:
    tmp = line.strip()
    if cogcat.match(tmp):
        pattern = cogcat.match(tmp)
        cogcatdict[pattern.group(2)] = pattern.group(1)[0]

cogcatfile.close()

glist = []

genome1x = [g1spanx1, g1spanx2]
genome1y = [2, 2]
genome2x = [g1spanx1, g1spanx2]
genome2y = [10, 10]
mid1x = [g1spanx1+200, g1spanx2-200]
mid1y = [10.75, 10.75]
mid2x = [g1spanx1+200, g1spanx2-200]
mid2y = [4.75, 4.75]

xdiff = 0

if g1spanx1 > g2spanx1:
    xdiff = g1spanx1 - g2spanx1
    padding = xdiff - g1starts[0]
else:
    xdiff = g2spanx1 - g1spanx1
    padding = g2starts[0] - xdiff

##Start plotting
fig = plt.figure(1, figsize=(16,5))
#ax1 = fig.add_subplot(211) #makes the subplot and squeezes the figure to half panel
ax1 = fig.add_subplot(111) #makes the full figure plot. larger.

ax1.plot(mid1x, mid1y, color='#CDAA7D', marker='|', mec='#CDAA7D', ls ='-', lw=1.0)
ax1.text(g1spanx1+200, 5.6, g1spanx1+200, fontsize=8, color='black', rotation=90)
ax1.text(g1spanx2-200, 5.6, g1spanx2-200, fontsize=8, color='black', rotation=90)
ax1.plot(mid2x, mid2y, color='#CDAA7D', marker='|', mec='#CDAA7D', ls ='-', lw=1.0)
ax1.text(g1spanx1+200, 11.6, g2spanx1+200, fontsize=8, color='black', rotation=90)
ax1.text(g1spanx2-200, 11.6, g2spanx2-200, fontsize=8, color='black', rotation=90)

#ax1.axis([0, g1length, 0, 14])
ax1.axis([g1spanx1, g1spanx2, 0, 16])

for k, v in g1featdict.iteritems():
    g1feat = v
    g1start = g1feat.location._start.position
    g1stop = g1feat.location._end.position
    g1size = g1stop - g1start + 1
    g1mid = g1start + ((g1stop - g1start) / 2.0)
    g1desc = g1feat.qualifiers['product'][0]
    g1gene = ""
    if g1feat.qualifiers.has_key('gene'):
        g1gene = g1feat.qualifiers['gene'][0] #displays gene name
        #g1gene = g1desc #displays product description
    else:
        g1gene = g1feat.qualifiers['locus_tag'][0] #displays locus tag
        #g1gene = g1desc #displays product description
    cogcolor = "#D3D3D3" #base color
    if g1feat.qualifiers.has_key('note'):
        cog = g1feat.qualifiers['note'][0]
        if cog in cogcatdict:
            cogcolor = cogdict[cogcatdict[cog]]
    if g1feat.type == 'tRNA':
        cogcolor = '#800000'
    if g1feat.type == 'rRNA':
        cogcolor = '#9400D3'
        g1gene = g1desc
    if g1feat.strand == -1:
        if g1spanx1 <= g1start <= g1spanx2 and g1spanx1 <= g1stop <= g1spanx2:
            rect = Rectangle((g1start, 4.0), g1size, 0.5, fc=cogcolor, 
                ec='#CDAA7D', alpha=0.5)
            plt.gca().add_patch(rect)
            #ax1.text(g1mid, 4.5, g1gene, fontsize=8, color='black', rotation=45)
            ax1.text(g1mid, 3.4, g1gene, fontsize=8, color='black', 
                horizontalalignment='center')
    else:
        if g1spanx1 <= g1start <= g1spanx2 and g1spanx1 <= g1stop <= g1spanx2:
            rect = Rectangle((g1start, 5.0), g1size, 0.5, fc=cogcolor, 
                ec='#CDAA7D', alpha=0.5)
            plt.gca().add_patch(rect)
            #ax1.text(g1mid, 5.5, g1gene, fontsize=8, color='black', rotation=45)
            ax1.text(g1mid, 5.6, g1gene, fontsize=8, color='black', 
                horizontalalignment='center')

#Genome 2
for k, v in g2featdict.iteritems():
    g2feat = v
    g2locustag = g2feat.qualifiers['locus_tag'][0]
    g2start = g2feat.location._start.position
    g2stop = g2feat.location._end.position
    g2size = g2stop - g2start + 1
    g2mid = g2start + ((g2stop - g2start) / 2.0)
    g2desc = g2feat.qualifiers['product'][0]
    g2gene = ""
    if g2feat.qualifiers.has_key('gene'):
        g2gene = g2feat.qualifiers['gene'][0] #displays gene name
        #g1gene = g1desc #displays product description
    else:
        g2gene = g2feat.qualifiers['locus_tag'][0] #displays locus tag
        #g1gene = g1desc #displays product description
    cogcolor = "#D3D3D3" #base color
    if g2locustag in g2cogs:
        if g2cogs[g2locustag] != '-':
            #cogcolor = cogdict[g2cogs[g2locustag]]
            cogcolor = cogdict[cogcatdict[g2cogs[g2locustag]]]
    if g2feat.type == 'tRNA':
        cogcolor = '#800000'
    if g2feat.type == 'rRNA':
        cogcolor = '#9400D3'
        g2gene = g2desc
    if g2feat.strand == -1:
        #if g1spanx1 <= g1start <= g1spanx2 and g1spanx1 <= g1stop <= g1spanx2:
        #if g2start >= g1spanx1 and g2stop <= g1spanx2:
        newg2start = g2start + padding
        newg2stop = g2stop + padding
        newg2mid = newg2start + ((newg2stop - newg2start) / 2.0)
        rect = Rectangle((newg2start, 10.0), g2size, 0.5, fc=cogcolor, 
            ec='#CDAA7D', alpha=0.5)
        plt.gca().add_patch(rect)
        #ax1.text(g2mid, 10.5, g2gene, fontsize=8, color='black', rotation=45)
        ax1.text(newg2mid, 9.4, g2gene, fontsize=8, color='black', 
            horizontalalignment='center')
        #print "g2start, newg2start", g2start, newg2start
    else:
        #if spanx1 <= g2start <= spanx2 and spanx1 <= g2stop <= spanx2:
        #if g2start >= g1spanx1 and g2stop <= g1spanx2:
        newg2start = g2start + padding
        newg2stop = g2stop + padding
        newg2mid = newg2start + ((newg2stop - newg2start) / 2.0)
        rect = Rectangle((newg2start, 11.0), g2size, 0.5, fc=cogcolor, 
            ec='#CDAA7D', alpha=0.5)
        plt.gca().add_patch(rect)
        #ax1.text(newg2mid, 11.5, g2gene, fontsize=8, color='black', rotation=45)
        ax1.text(newg2mid, 11.6, g2gene, fontsize=8, color='black', 
            horizontalalignment='center')
        #print "g2start, newg2start", g2start, newg2start


#Draw legend box for COG categories
coglist = []
for k, v in cogdict.iteritems():
    coglist.append((k,v))

coglist.sort()

ccounts = len(coglist)
a = 0
spansize = g1spanx2 - g1spanx1
spanmid = g1spanx1 + ((g1spanx2 - g1spanx1) / 2.0)
xstart = spanmid
#increment = 20000 #for synecoccus
increment = spansize * 0.01 #for gloeobacter
while a < ccounts:
    fc = coglist[a][1]
    tx = coglist[a][0]
    #rect = Rectangle((xstart, 2.5), 20000, 0.25, facecolor=fc, alpha=0.5) #for synecococcus
    rect = Rectangle((xstart, 2.5), increment, 0.5, fc=fc, alpha=0.5) #for gloeobacter
    plt.gca().add_patch(rect)
    ax1.annotate(tx, xy=(xstart+(increment/2.0), 2.2), 
        horizontalalignment='center', verticalalignment='center', fontsize=8)
    a += 1
    xstart = xstart + increment

ax1.annotate('COG categories', xy=(spanmid, 1.2), horizontalalignment='left', 
    verticalalignment='center', fontsize=10)
ax1.annotate(g1seq.annotations['organism'], xy=(0.1, 0.1), 
    xycoords='axes fraction', horizontalalignment='left', 
    verticalalignment='center', fontsize=10)
ax1.annotate(g2seq.annotations['organism'], xy=(0.1, 0.9), 
    xycoords='axes fraction', horizontalalignment='left', 
    verticalalignment='center', fontsize=10)

frame1 = plt.gca()
for tick in frame1.axes.get_yticklines():
    tick.set_visible(False)
for y in frame1.axes.get_yticklabels():
    y.set_visible(False)
ax1.grid(False)


plt.show()


